package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"
)

// ForkRun 处理 wkhtmltopdf 命令的执行
func ForkRun(ch chan string) {
	logger := log.New(os.Stdout, "Image: ", log.LstdFlags)
	for {
		args := <-ch
		logger.Printf("Received: %v", args)

		as := strings.Split(args, ",")
		logger.Printf("Command= xvfb-run %s", strings.Join(as, " "))

		cmd := exec.Command("xvfb-run", as...)
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			logger.Printf("Error creating stdout pipe: %v", err)
			ch <- "ERROR"
			continue
		}

		if err := cmd.Start(); err != nil {
			logger.Printf("Error starting command: %v", err)
			ch <- "ERROR"
			continue
		}

		success := true
		scanner := bufio.NewScanner(stdout)
		for scanner.Scan() {
			line := scanner.Text()
			logger.Printf("STDOUT: %s", line)
			if strings.Contains(line, "Done") {
				ch <- "DONE"
				break
			} else if strings.Contains(line, "<Error>") {
				success = false
				cmd.Process.Kill()
				ch <- "FATAL"
				break
			}
		}

		if err := scanner.Err(); err != nil {
			logger.Printf("Error reading stdout: %v", err)
			success = false
			ch <- "ERROR"
		}

		if err := cmd.Wait(); err != nil {
			logger.Printf("Error waiting for command: %v", err)
			success = false
			ch <- "ERROR"
		}

		if !success {
			continue
		}
		logger.Printf("Done: %v", args)
	}
}

// MakeCommand 根据请求生成 wkhtmltopdf 命令
func MakeCommand(rb RequestBody) string {
	cmd := ""
	if rb.ToType == 1 {
		cmd = "wkhtmltopdf,--minimum-font-size,1,--encoding,utf-8"
	} else {
		cmd = "wkhtmltoimage"
	}

	for k, v := range rb.Args {
		if v == "" {
			continue
		}
		cmd += fmt.Sprintf(",--%s,%s", k, v)
	}
	cmd += fmt.Sprintf(",%s,%s", rb.HtmlPath, rb.extPath)
	return cmd
}

// NewWorker 创建一个新的 Worker
func NewWorker() chan string {
	ch := make(chan string)
	go ForkRun(ch)
	return ch
}

// WorkerPool 管理 Worker
type WorkerPool struct {
	lock   sync.Mutex
	avail  []chan string
	logger *log.Logger
}

// NewWorkerPool 创建一个新的 WorkerPool
func NewWorkerPool() *WorkerPool {
	wp := new(WorkerPool)
	wp.avail = append(wp.avail, NewWorker())
	wp.logger = log.New(os.Stdout, "POOL: ", log.LstdFlags)
	return wp
}

// RequestWorker 从池中获取一个 Worker
func (wp *WorkerPool) RequestWorker() chan string {
	wp.lock.Lock()
	defer wp.lock.Unlock()
	wp.logger.Printf("Requesting: Avail=%d", len(wp.avail))
	if len(wp.avail) > 0 {
		x, avail := wp.avail[len(wp.avail)-1], wp.avail[:len(wp.avail)-1]
		wp.avail = avail
		return x
	}
	return NewWorker()
}

// Recycle 回收 Worker
func (wp *WorkerPool) Recycle(ch chan string) {
	wp.lock.Lock()
	defer wp.lock.Unlock()
	wp.avail = append(wp.avail, ch)
	wp.logger.Printf("Recycled: Avail=%d", len(wp.avail))
}

// RequestBody 定义请求体结构
type RequestBody struct {
	Html     string            `json:"html"`
	Args     map[string]string `json:"args"`
	Ext      string            `json:"ext"`
	ToType   int               `json:"toType"` // 转换类型：1：pdf 2: image
	HtmlPath string            // HTML 文件路径
	Url      string            `json:"url"` // 转 URL 为 PDF/Image
	extPath  string
}

// parseRequest 解析请求体
func parseRequest(r *http.Request, dst interface{}) error {
	if !strings.Contains(r.Header.Get("Content-Type"), "application/json") {
		return fmt.Errorf("Content-Type header must be application/json")
	}
	body, err := io.ReadAll(r.Body)
	if err != nil {
		return err
	}
	defer r.Body.Close()
	return json.Unmarshal(body, dst)
}

// HtmlToPdfHandler 处理 HTTP 请求
func HtmlToPdfHandler(pool *WorkerPool, logger *log.Logger) http.HandlerFunc {
	return func(resp http.ResponseWriter, req *http.Request) {
		logger.Printf("RequestURI: %v", req.RequestURI)

		var htmlfilename string
		var rb RequestBody

		switch req.Method {
		case "POST":
			html, err := os.CreateTemp("/tmp", "html")
			if err != nil {
				logger.Printf("Error creating temp file: %v", err)
				resp.WriteHeader(http.StatusInternalServerError)
				return
			}
			defer os.Remove(html.Name())

			err = parseRequest(req, &rb)
			if err != nil {
				logger.Printf("Error parsing request: %v", err)
				http.Error(resp, err.Error(), http.StatusInternalServerError)
				return
			}

			if _, err := io.WriteString(html, rb.Html); err != nil {
				logger.Printf("Error writing to temp file: %v", err)
				resp.WriteHeader(http.StatusInternalServerError)
				return
			}
			htmlfilename = html.Name() + ".html"
			rb.HtmlPath = htmlfilename
			if rb.ToType == 0 {
				rb.ToType = 1
			}
			if err := os.Rename(html.Name(), htmlfilename); err != nil {
				logger.Printf("Error renaming temp file: %v", err)
				resp.WriteHeader(http.StatusInternalServerError)
				return
			}
			defer os.Remove(htmlfilename)
		case "GET":
			htmlfilename = req.FormValue("url")
			if htmlfilename == "" {
				logger.Printf("GET: url not given")
				resp.WriteHeader(http.StatusBadRequest)
				return
			}

			queryParams := req.URL.Query()
			paramsMap := make(map[string]string)
			for key, values := range queryParams {
				if key == "toType" || key == "url" {
					continue
				}
				if len(values) > 0 {
					paramsMap[key] = values[0]
				}
			}

			toTypeStr := req.FormValue("toType")
			if toTypeStr == "" {
				toTypeStr = "1"
			}
			toType, _ := strconv.Atoi(toTypeStr)
			rb = RequestBody{
				HtmlPath: htmlfilename,
				Url:      htmlfilename,
				ToType:   toType,
				Args:     paramsMap,
			}
		default:
			logger.Printf("Method not supported: %s", req.Method)
			resp.WriteHeader(http.StatusMethodNotAllowed)
			return
		}

		pattern := "pdf"
		if rb.ToType != 1 {
			pattern = "jpg"
		}
		extFile, err := os.CreateTemp("/tmp", pattern)
		if err != nil {
			logger.Printf("Error creating temp file: %v", err)
			resp.WriteHeader(http.StatusInternalServerError)
			return
		}
		defer os.Remove(extFile.Name())

		rb.extPath = extFile.Name()
		if rb.ToType == 1 {
			rb.extPath += ".pdf"
		} else {
			rb.extPath += ".jpg"
		}
		if err := os.Rename(extFile.Name(), rb.extPath); err != nil {
			logger.Printf("Error renaming temp file: %v", err)
			resp.WriteHeader(http.StatusInternalServerError)
			return
		}
		defer os.Remove(rb.extPath)

		retry := 3
		success := false
		for retry > 0 {
			ch := pool.RequestWorker()
			ch <- MakeCommand(rb)
			flag := <-ch
			if flag == "DONE" {
				if rb.ToType == 1 {
					resp.Header().Set("Content-Type", "application/pdf")
				} else {
					resp.Header().Set("Content-Type", "image/jpeg")
				}
				file, err := os.Open(rb.extPath)
				if err != nil {
					logger.Printf("Error opening output file: %v", err)
					resp.WriteHeader(http.StatusInternalServerError)
					return
				}
				defer file.Close()

				buf := make([]byte, 32*1024) // 32KB buffer
				if _, err := io.CopyBuffer(resp, file, buf); err != nil {
					logger.Printf("Error copying file to response: %v", err)
					resp.WriteHeader(http.StatusInternalServerError)
					return
				}
				pool.Recycle(ch)
				success = true
				break
			} else if flag == "FATAL" {
				logger.Printf("Fatal Error! No Retry")
				break
			}
			retry--
			logger.Printf("Error! Retry Left: %d", retry)
		}

		if !success {
			logger.Printf("Conversion FAILED!")
			resp.WriteHeader(http.StatusInternalServerError)
		}
	}
}

func main() {
	mux := http.NewServeMux()
	pool := NewWorkerPool()
	logger := log.New(os.Stdout, "HTTP: ", log.LstdFlags)
	mux.HandleFunc("/", HtmlToPdfHandler(pool, logger))

	server := &http.Server{
		Addr:         ":8080",
		Handler:      mux,
		ReadTimeout:  60 * time.Second,
		WriteTimeout: 60 * time.Second,
	}
	logger.Printf("Server started at http://localhost:8080")
	if err := server.ListenAndServe(); err != nil {
		logger.Fatalf("Server error: %v", err)
	}
}
